﻿from PyQt5.QtWidgets import (QApplication, QWidget, QPushButton, QVBoxLayout, QHBoxLayout, QLabel, QFileDialog, QTextEdit)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap
from PyQt5.QtWebEngineWidgets import QWebEngineView
from PyQt5.QtCore import QUrl
import sys
import folium
from picarta import Picarta
import os
import tempfile

class GeoLocalizationApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        
        # 初始化 Picarta API
        self.api_token = "64J0BIKWGANTAAJYDGNS"
        self.localizer = Picarta(self.api_token)
        self.result = None
        
    def initUI(self):
        # 设置窗口属性
        self.setWindowTitle('街景地理定位')
        self.setGeometry(200, 200, 1000, 750)
        
        # 创建布局
        self.main_layout = QVBoxLayout()
        self.top_layout = QHBoxLayout()
        self.middle_layout = QHBoxLayout()
        self.bottom_layout = QHBoxLayout()
        
        # 创建控件
        self.selectButton = QPushButton('选择图片', self)
        self.selectButton.clicked.connect(self.selectImage)
        
        self.localizationButton = QPushButton('地理定位', self)
        self.localizationButton.clicked.connect(self.performLocalization)
        self.localizationButton.setEnabled(False)  # 初始禁用

        self.imageLabel = QLabel(self)
        self.imageLabel.setAlignment(Qt.AlignCenter)
        self.imageLabel.setFixedSize(400, 300)
        self.imageLabel.setScaledContents(True)
        self.imageLabel.setStyleSheet("background-color: #f0f0f0; border: 1px solid #cccccc;")

        self.resultText = QTextEdit(self)
        self.resultText.setReadOnly(True)
        self.resultMiddleText = QTextEdit(self)
        self.resultMiddleText.setReadOnly(True)
        
        # 地图视图
        self.mapView = QWebEngineView(self)
        self.mapView.setFixedSize(400, 300)
        
        # 添加控件到布局
        self.top_layout.addWidget(self.selectButton)
        self.top_layout.addWidget(self.localizationButton)
        
        self.middle_layout.addWidget(self.imageLabel)
        self.middle_layout.addWidget(self.resultMiddleText)

        self.bottom_layout.addWidget(self.mapView)
        self.bottom_layout.addWidget(self.resultText)

        self.main_layout.addLayout(self.top_layout)
        self.main_layout.addLayout(self.middle_layout)
        self.main_layout.addLayout(self.bottom_layout)
        
        # 设置布局
        self.setLayout(self.main_layout)
        
    def selectImage(self):
        # 打开文件对话框，选择图片
        options = QFileDialog.Options()
        options |= QFileDialog.ReadOnly
        self.img_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "图片文件 (*.jpg *.jpeg *.png)", options=options)
        
        if self.img_path:
            self.localizationButton.setEnabled(True)  # 启用地理定位按钮
            self.resultMiddleText.clear()
            self.resultText.clear()  # 清空结果展示区
            # 显示选中的图片
            pixmap = QPixmap(self.img_path)
            self.imageLabel.setPixmap(pixmap)
    
    def performLocalization(self):
        if hasattr(self, 'img_path') and self.img_path:
            # 调用 Picarta API 进行地理定位
            self.result = self.localizer.localize(img_path=self.img_path)
            self.displayResult(self.result)
            self.displayMap(self.result)
    
    def displayResult(self, result):
        if result:
            # 清空文本框并开始构建结果展示
            self.resultText.clear()
            self.resultMiddleText.clear()
            
            # 打印基本信息
            result_middle_str = f"AI 置信度: {result.get('ai_confidence', 'N/A')}\n"
            result_middle_str += f"国家: {result.get('ai_country', 'N/A')}\n"
            result_middle_str += f"城市: {result.get('city', 'N/A')}\n"
            result_middle_str += f"省份/州: {result.get('province', 'N/A')}\n"
            result_middle_str += f"纬度: {result.get('ai_lat', 'N/A')}\n"
            result_middle_str += f"经度: {result.get('ai_lon', 'N/A')}\n\n"
            
            result_str = "前 10 个预测结果：\n"
            for i in range(1, 11):
                prediction = result.get('topk_predictions_dict', {}).get(str(i), {})
                address = prediction.get('address', {})
                city = address.get('city', 'N/A')
                province = address.get('province', 'N/A')
                country = address.get('country', 'N/A')
                
                result_str += f"{i}. 地址: {city}, {province}, {country}\n"
                result_str += f"   置信度: {prediction.get('confidence', 'N/A')}\n"
                result_str += f"   GPS 坐标: {prediction.get('gps', 'N/A')}\n\n"
            
            result_str += "=" * 50
            
            self.resultMiddleText.setText(result_middle_str)
            self.resultText.setText(result_str)
        else:
            self.resultText.setText("定位失败或未获取到结果")
    
    def displayMap(self, result):
        if result:
            m = folium.Map(location=[20, 0], zoom_start=2)
            
            for i in range(1, 11):
                prediction = result.get('topk_predictions_dict', {}).get(str(i), {})
                gps = prediction.get('gps', 'N/A')
                
                if isinstance(gps, list) and len(gps) == 2:
                    lat, lon = gps
                    folium.Marker([lat, lon], popup=f'预测 {i}').add_to(m)
                elif isinstance(gps, str):
                    try:
                        lat, lon = map(float, gps.split(','))
                        folium.Marker([lat, lon], popup=f'预测 {i}').add_to(m)
                    except:
                        print(f"无法解析 GPS 坐标: {gps}")
                else:
                    print(f"无效的 GPS 坐标格式: {gps}")
            
            # 保存地图到临时文件
            temp_dir = tempfile.gettempdir()
            map_file = os.path.join(temp_dir, 'map.html')
            m.save(map_file)
            
            if os.path.exists(map_file):
                # 显示地图
                self.mapView.load(QUrl.fromLocalFile(map_file))
                print(f"地图文件已保存到: {map_file}")
            else:
                print("地图文件保存失败")
        else:
            # 如果没有坐标信息，显示默认地图
            m = folium.Map(location=[20, 0], zoom_start=2)
            temp_dir = tempfile.gettempdir()
            map_file = os.path.join(temp_dir, 'map.html')
            m.save(map_file)
            
            # 检查文件是否正确保存
            if os.path.exists(map_file):
                self.mapView.load(QUrl.fromLocalFile(map_file))
                print(f"默认地图文件已保存到: {map_file}")
            else:
                print("默认地图文件保存失败")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    from PyQt5.QtWebEngineWidgets import QWebEngineView
    ex = GeoLocalizationApp()
    ex.show()
    sys.exit(app.exec_())